package com.example.ai.controller;

import com.example.ai.entity.AIRequest;
import com.example.ai.service.AIService;
import jakarta.validation.Valid;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.time.Duration;

/**
 * AI controller
 *
 * @author liujian
 */
@RestController
@RequestMapping("/ai")
public class AIController {

    /**
     * 日志
     */
    private static final Logger logger = LoggerFactory.getLogger(AIController.class);

    /**
     * AI service
     */
    @Autowired
    private AIService aiService;

    /**
     * 对话
     *
     * @param request  请求参数
     * @param provider 大模型名称
     * @return 结果
     */
    @PostMapping("/chat")
    public ResponseEntity<String> generateText(
            @Valid @RequestBody AIRequest request,
            @RequestParam(defaultValue = "qianwen") String provider
    ) {
        return ResponseEntity.ok(aiService.generateText(request, provider));
    }

    /**
     * 流式返回
     *
     * @param request  请求参数
     * @param provider 大模型名称
     * @return 结果
     */
    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamGenerate(
            @Valid AIRequest request,
            @RequestParam(defaultValue = "qianwen") String provider,
            @RequestHeader(value = "X-User-Id", defaultValue = "crm") String userId) {
        return aiService.streamGenerate(request, provider, userId)
                .onBackpressureBuffer(50) // 背压控制
                .timeout(Duration.ofMinutes(5)) // 超时设置
                .doOnCancel(() -> logger.info("客户端断开连接"));
    }

}
